#ifndef STATSxx_MACHINE_LEARNING_NEURAL_NETWORK_DEEP_BELIEF_NETWORK_DBN_HPP
#define STATSxx_MACHINE_LEARNING_NEURAL_NETWORK_DEEP_BELIEF_NETWORK_DBN_HPP


// STL
#include <ostream> // std::ostream
#include <vector>  // std::vector<>

// jScience
#include "jScience/linalg/Matrix.hpp" // Matrix<>
#include "jScience/linalg/Vector.hpp" // Vector<>
#include "jScience/stl/ostream.hpp"   // NullStream

// stats++
#include "statsxx/machine_learning/NeuralNet.hpp"                        // NEURAL_NET
#include "statsxx/machine_learning/restricted_Boltzmann_machine/RBM.hpp" // RBM


namespace neural_network
{
    //=========================================================
    // DEEP BELIEF NETWORKS
    //=========================================================

    // TODO: need to read the Hinton Science dimensionality-reduction paper ... it remarks on weights needing to be initialized in just the right way (must make sure that the way I do it is consistent / correct)

    class DBN
    {

    public:

        // << jmm: not sure if this should be part of this or RBM >>
        enum class PropagationType
        {
            stochastic,        // fully stochastic sampling
            pseudo_stochastic, // stochastic sampling only in the top layer
            deterministic      // fully deterministic sampling
        };


        // basic constructor for a DBN
        DBN(const std::vector<int> &arch);
        ~DBN();


        // greedy, layer-by-layer unsupervised training of the RBM
        void pretrain(
                      const int ntrials,                // number of trials to use for each RBM initialization
                      // -----
                      const PropagationType prop_type,  // type of propagation of signal to use in greedy training
                      // -----
                      // the parameters below are identical to RBM::stochastic_gradient():
                      const int max_epochs,             // maximum number of epochs
                      const int nbatches,               // number of batches per epoch
                      // -----
                      double lr,                        // learning rate
                      double momentum,                  // momentum
                      // -----
                      const double weight_penalty,      // weight penalty
                      // -----
                      const double sparsity_target,     // sparsity target
                      const double sparsity_penalty,    // sparsity penalty
                      const double sparsity_multiplier, // multiplier for q (hidden activation) EMA
                      const double q_min,               // minimum (desired) q
                      const double q_max,               // maximum (desired) q
                      const double q_penalty,           // q penalty
                      // -----
                      const int K_begin,                // number of starting Monte Carlo (MC) itertions
                      const int K_end,                  // number of ending Monte Carlo itertions
                      const double K_rate,              // rate of updates to number of MC iterations
                      // -----
                      const bool mean_field,            // mean-field approximation
                      // -----
                      const double convg_criterion,     // convergence criterion
                      const int max_no_improvement,     // max number of epochs to go without an improvement in convergence
                      // -----
                      const Matrix<double> &X,          // data
                      const bool x_binomial,            // whether the data is binomial
                      // -----
                      std::ostream* os = &NullStream    // (optional) output stream
                      );


        // << jmm: this finetune is going to be come the up-down supervised thing -- see Hinton >>
        // share the weights of the DBN with a MLP, and finetune them using backpropagation
        void finetune(
                      const int imethod,                // type of method to use for supervised MLP training
                      // -----
                      const Matrix<double> &X,          // data
                      const Matrix<double> &X_out,      // data (output)
                      const bool x_out_classif
                      );

        // << jmm: possibly deprecated for prop_forward() and prop_backward() ... need to think ... maybe will be removed>>
        //
        // get the activations (of a layer or the whole network)
        std::vector<Vector<double>> activations(
                                                const PropagationType prop_type,  // type of propagation of signal
                                                // -----
                                                const Vector<double> &x           // data
                                                ) const;

        Vector<double> activations(
                                   const int layer,                  // layer to propagate signal to
                                   // -----
                                   const PropagationType prop_type,  // type of propagation of signal
                                   // -----
                                   const Vector<double> &x           // data
                                   ) const;


        // generate a sample from the distribution of patterns learned (generative sampling)
        Vector<double> sample(
                              const bool mean_field,                // mean-field approximation
                              // -----
                              const int K,                          // number of Markov iterations
                              // -----
                              Vector<double> x = Vector<double>()   // initial vector (optional -- otherwise, a random activation vector at the top RBM is generated)
                              ) const;

        // predict
        // << jmm: should be part of DBN+GP or DBN+MLP, NOT here >>
        Vector<double> predict(
                               const Vector<double> &input
                               );

        std::vector<int> get_architecture() const;

        // MLP
        NEURAL_NET MLP;

        // TODO: I'm not sure how this all should work out ... not sure we want this public ... comments have previously been made though about integrating the NN into here
        std::pair<
                  Matrix<double>,
                  Vector<double>
                  > get_RBM_weights();

    private:

        // ARCHITECTURE, ETC.
        std::vector<int> architecture;

        // RBMS
        std::vector<restricted_Boltzmann_machine::RBM> RBM;

        // MLP
//        NEURAL_NET MLP;

        Vector<double> prop_forward(
                                    const int layer0,                  // layer to propagate signal from
                                    const int layern,                  // "                       " to
                                    // -----
                                    const PropagationType prop_type,   // type of propagation of signal
                                    // -----
                                    Vector<double> x                   // data
                                    ) const;

        Vector<double> prop_back(
                                 const int layer0,                     // layer to propagate signal from
                                 const int layern,                     // "                       " to
                                 // -----
                                 const PropagationType prop_type,      // type of propagation of signal
                                 // -----
                                 Vector<double> x                      // data
                                 ) const;

        // << jmm: deprecated for activations() ... which in turn is possibly deprecated for prop_forward() and prop_backward() ... will be removed>>
        Vector<double> RBM_propagate(
                                     const PropagationType prop_type,  // type of propagation of signal
                                     // -----
                                     const int layer,                  // layer to propagate signal to
                                     // -----
                                     const Vector<double> &x           // data
                                     ) const;
        /*
        std::pair<
                  Matrix<double>,
                  Vector<double>
                  > get_RBM_weights();
        */

    };
}

#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_pretrain.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_finetune.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_sample.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_predict.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_get_architecture.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/activations.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_RBM_propagate.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_prop_forward.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_prop_back.cpp"
#include "statsxx/machine_learning/neural_network/deep_belief_network/src/DBN_get_RBM_weights.cpp"


#endif
